In [1]:
import scipy.integrate
import matplotlib.pyplot as plt
import numpy as np
%matplotlib inline

In [2]:
g = 9.8
k = 100000
def y(t, x):
    if x[0] <0:
        a = -g -k*x[0]
    else:
        a = -g
    xdot = [x[1], a]
    return xdot
ode = scipy.integrate.ode(y)
ode.set_integrator('dopri5', rtol=1e-10)
ode.set_initial_value([1,0], 0)
y = []
t = []
dt = 0.01
tf = 10
n_t = int(tf/dt)
y = np.zeros((n_t,2))
t = np.zeros(n_t)
count = 0
while ode.t + dt< tf:
    ode.integrate(ode.t + dt)
    y[count] = ode.y
    t[count] = ode.t
    count = count + 1
t = t[:count]
y = y[:count]
plt.plot(t,y[:,0])


Out[2]:
[<matplotlib.lines.Line2D at 0x4ae2790>]

Calculate the energy to show how much the numerical scheme violates the conversation of energy.


In [3]:
x = y[:,0]
x_dot = y[:,1]
e = x_dot**2/2 + 1*9.8*x + np.where(x < 0, 1, 0)*k*x**2/2

This numerical integration scheme is fairly accrurate because the relative tolerance ensures that the step size is reduced during the discontinuity in acceleration. However, during each bounce there is still a slight change in energy.


In [4]:
plt.plot(t, e)


Out[4]:
[<matplotlib.lines.Line2D at 0x4d0cad0>]